# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: jhkim

import os
import datetime
import fire
import pandas
import cc.appcoin2.strategy.support.interval_feed_config_gen as ifcg
from xunkemgmt_client.tool.slack_noti import send_to_slack

def sbatch_cmds(*, cmds, ncpu, driver_dir, job_name, tsstr, dry):
  job_filename = f'{driver_dir}/sbatch_{job_name}_{tsstr}.sh'
  sbatch_out_file = os.path.realpath(f'sbatch_out.{tsstr}.txt')
  with open(job_filename, 'w') as job_file:
    wrapcmds = (
        ["#!/bin/bash"]
        + [f"#SBATCH -o {sbatch_out_file} "]
        + [""]
        + [f"{cmd} &" for cmd in cmds] + ["wait"])
    jobstr = "\n".join(wrapcmds)
    job_file.write(jobstr)
  print("# of jobs: %s" % len(cmds))
  print(job_filename)
  print(sbatch_out_file)
  cmd_out = f"sbatch -n {ncpu} --job-name={job_name} -c 1 --mem-per-cpu=4gb --wait {job_filename}"
  print(cmd_out)
  if dry:
    pass
  else:
    os.system(cmd_out)


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]

def gen_interval_feed(
    trading_date,
    ncpu,
    dry,
    resolution_name,
    interval_feed_machine="main",
    mea_equals='',
    mea_contains='',
    worker='',
    regen_only=False):
  resolution = {'PT1M': 60, 'P1DT': 86400}[resolution_name]
  interval_feed_root = "/remote/iosg/data-2/buckets/feed.derived.interval_h5"
  out_cache_root = f"{interval_feed_root}/coin/{interval_feed_machine}/{resolution_name}"
  tsstr = f'{datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")}'
  job_name = f"regen_cache_symbol_list_{resolution_name}_{tsstr}"
  driver_dir = os.path.join("regen_cache", job_name)
  os.makedirs(driver_dir, exist_ok=True)

  tds = get_trading_dates(trading_date)
  cmds = []
  mea_list= set()
  for td in tds:
    for dir_name, _, file_names in os.walk("python/coin_deploy/feed/interval_feed"):
      for file_name in file_names:
        tdstr = td.strftime('%Y%m%d')
        tddir = os.path.join(driver_dir, tdstr)
        os.makedirs(tddir, exist_ok=True)
        meas = ifcg.gen_config(
            os.path.join(dir_name, file_name),
            output_dir=f"{tddir}",
            pi_root='./data/coin2/product_info',
            feed_root='/remote/iosg/coin-2/feed_cache2',
            mea_equals=mea_equals,
            mea_contains=mea_contains,
            worker=worker,
            date=tdstr)
        for mea in meas:
          mea_list.add(mea)
          driver_file = os.path.join(tddir, mea + '.json')
          tdmeadir = os.path.join(out_cache_root, mea, tdstr)
          tdmeatime = os.path.join(tdmeadir, f"{mea}--time.h5")
          if regen_only and os.path.exists(tdmeadir) and os.path.exists(tdmeatime):
            os.remove(driver_file)
            continue
          cmd = f"""srun -N1 -n1 --exclusive ../coin_binary/coin2-bazel-bin-3.10/cc/appcoin2/support/feed/interval_feed_tool/interval_feed_generator \
--resolution={resolution} \
--driver={driver_file} \
--trading_date={td.strftime('%Y%m%d')} \
--enable_time=True \
--interval_feed_root={interval_feed_root} \
--interval_feed_machine={interval_feed_machine}"""
          if mea == "Spot.Coingecko.v3":
            cmd += f""" --enable_ohlc=false \
            --enable_vwap=false \
            --enable_volume=false \
            --enable_other_stats=false \
            --enable_time=false \
            --enable_market_analysis=true """
          cmds.append(cmd)
          if len(cmds) % 15 == 0:
            cmds.append("sleep 0.5")
  def _chunks(l, m):
    for i in range(0, len(l), m):
      yield l[i:i + m]
  for subcmds in _chunks(cmds, ncpu * 2):
    sbatch_cmds(cmds=subcmds, ncpu=ncpu, driver_dir=driver_dir, job_name=job_name, tsstr=tsstr, dry=dry)
  verify_h5_file(mea_list,out_cache_root,tds,resolution)


def verify_h5_file(mea_list,file_path,tds,resolution):
  slack_msg=[f"[interval feed verify]{resolution} interval feed missing : \n",]
  for td in tds:
    td=td.strftime("%Y%m%d")
    for mea in mea_list:
      if mea in ["Futures.Ftx.v1", "Spot.Ftx.v1"]:
        continue
      check_path = os.path.join(file_path, mea, td)
      if mea == "Spot.Coingecko.v3":
        check_path = os.path.join(check_path, "Spot.Coingecko.v3--market_analysis.h5")
        if not os.path.exists(check_path) or os.path.getsize(check_path) == 0:
          slack_msg.append(f"    mea:{mea}        trading date:{td} \n")
      elif not os.path.exists(check_path) or len(os.listdir(check_path)) != 5:
        slack_msg.append(f"    mea:{mea}        trading date:{td} \n")
      else:
        for file in os.listdir(check_path):
          if os.path.getsize(os.path.join(check_path,file)) == 0 :
            slack_msg.append(f"    mea:{mea}        trading date:{td} \n")
            break
  if len(slack_msg) > 1:
    slack_msg = ''.join(slack_msg)
    print(slack_msg)
    send_to_slack(slack_msg,'#coin_feed_noti', 'msg', mention_list = ['ziyan','leon'])
        
if __name__ == "__main__":
  fire.Fire(gen_interval_feed)
